from pygds import GDS
import pandas as pd

from project_setting import NEO4J_URI,NEO4j_AUTH




def get_neo4j_random_walks(save_path,neo4j_query_dict,steps=10,walks=200000,mode='node2vec',return_p=1,inOut_q=2):
    '''

    :param neo4j_query_dict: random_walk查询dict
    :param steps: 每一次随机游走的步长
    :param walks: 生成的总数
    :param mode: ‘random’ or ‘node2vec’
    :param return_p: 当前节点返回概率 BFS
    :param inOut_q: 当前节点深度搜索概率 DFS
    :return: 路径节点dataframe
    '''

    neo4j_query_dict['steps'] = steps
    neo4j_query_dict['walks'] = walks
    neo4j_query_dict['mode'] = mode
    neo4j_query_dict['return'] = return_p
    neo4j_query_dict['inOut'] = inOut_q

    # 参数链接 https://neo4j.com/docs/graph-data-science/current/alpha-algorithms/random-walk/
    with GDS(NEO4J_URI,auth=NEO4j_AUTH) as gds:
        result = gds.alpha.randomWalk.stream(
            neo4j_query_dict)

        _tmp = pd.DataFrame.from_records(result)
        _tmp.to_csv(save_path)
    return _tmp